package cn.iocoder.yudao.core.thread.core.utils;

import lombok.extern.slf4j.Slf4j;
import org.slf4j.MDC;

import java.util.Objects;
import java.util.UUID;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.Executor;
import java.util.function.Supplier;

@Slf4j
public class FutureUtils {

    private static Executor threadPoolExecutor;

    public static CompletableFuture<Void> runAsync(Runnable run) {

        return CompletableFuture.runAsync(run, threadPoolExecutor).handle((result, e) -> {
            if (Objects.nonNull(e)) {
                log.error(" 线程执行异常------>>> errMsg:{}", e.getMessage(), e);
                throw new RuntimeException(e);
            }
            return result;
        });
    }

    public static <U> CompletableFuture<U> supplyAsync(Supplier<U> supplier) {

        return CompletableFuture.supplyAsync(supplier, threadPoolExecutor).handle((result, e) -> {
            if (Objects.nonNull(e)) {
                log.error(" 线程执行异常------>>> errMsg:{}", e.getMessage(), e);
                throw new RuntimeException(e);
            }
            return result;
        });

    }

    public static <U> U get(CompletableFuture<U> completableFuture) {
        try {
            return completableFuture.get();
        } catch (Exception e) {
            log.error("从CompletableFuture获取数据异常 errMsg:{} error:{}", e.getMessage(), e);
        }
        return null;
    }

    public static void setThreadPoolExecutor(Executor executor) {
        threadPoolExecutor = executor;
//        if (executor instanceof TraceableExecutorWrapper) {
//            threadPoolExecutor = executor;
//        } else {
//            threadPoolExecutor = new TraceableExecutorWrapper(executor);
//        }

    }


    public static class TraceableExecutorWrapper implements Executor {
        Executor delegate;

        public TraceableExecutorWrapper(Executor delegate) {
            this.delegate = delegate;
        }

        @Override
        public void execute(Runnable command) {
            String traceId = MDC.get("__traceId");
            this.delegate.execute(new TraceRunnable(command, traceId));
        }
    }

    public static class TraceRunnable implements Runnable {
        private final Runnable delegate;
        private final String tracerId;

        public TraceRunnable(Runnable delegate, String tracerId) {
            this.delegate = delegate;
            this.tracerId = tracerId;
        }

        @Override
        public void run() {
            if (tracerId != null) {
                markChildTraceId(this.tracerId);
            }
            this.delegate.run();
        }

        /**
         * 用于多线程情况下
         *
         * @param traceId traceId
         */
        public void markChildTraceId(String traceId) {
            StringBuilder sb = new StringBuilder();
            String childTraceId = UUID.randomUUID().toString().replace("-", "");
            MDC.put("__traceId", sb.append(traceId).append(":").append(childTraceId).toString());
        }
    }

    /**
     * 阻塞，等待所有CompletableFuture执行完成
     */
    public static void join(CompletableFuture<?>... cfs) {
        CompletableFuture.allOf(cfs).join();
    }

}
